package com.cmz.common.utils;

import com.google.common.collect.Maps;

import java.util.HashMap;
import java.util.Map;

/**
 * 线程上下文工具
 * @version 1.0.0
 * @author cmz
 */
public class ThreadContextUtils {

    /** 线程上下文 */
    private static final ThreadLocal<Map<Object, Object>> CONTEXT = new InheritableThreadLocalMap();

    private ThreadContextUtils() {}

    /**
     * 获取线程上下文
     * @since 1.0.0
     * @return 返回线程上下文 Map 对象
     */
    public static Map<Object, Object> getResources() {
        return new HashMap<>(CONTEXT.get());
    }

    /**
     * 设置线程上下文
     * @since 1.0.0
     * @param contextMap 线程上下文 Map 对象
     */
    public static void setResources(Map<Object, Object> contextMap) {
        if (contextMap == null || contextMap.isEmpty()) {
            return;
        }

        CONTEXT.get().clear();
        CONTEXT.get().putAll(contextMap);
    }

    /**
     * 获取线程上下文属性
     * @since 1.0.0
     * @param key 属性 key
     * @return 返回线程上下文属性值
     */
    public static Object get(Object key) {
        return CONTEXT.get().get(key);
    }

    /**
     * 获取线程上下文属性
     * @since 2.0.0
     * @param key   属性 key
     * @param clazz 对象类型
     * @param <T>
     * @return 返回线程上下文属性值
     */
    public static <T> T get(Object key, Class<T> clazz) {
        Object object = CONTEXT.get().get(key);
        return object == null ? null : (T) object;
    }

    /**
     * 设置线程上下文属性
     * @since 1.0.0
     * @param key   属性 key
     * @param value 属性值
     */
    public static void put(Object key, Object value) {
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }

        if (value == null) {
            remove(key);
            return;
        }

        CONTEXT.get().put(key, value);
    }

    /**
     * 删除线程上下文属性
     * @since 1.0.0
     * @param key 属性 key
     * @return 线程上下文
     */
    public static Object remove(Object key) {
        return CONTEXT.get().remove(key);
    }

    /**
     * 删除线程上下文
     * @since 1.0.0
     */
    public static void remove() {
        CONTEXT.remove();
    }

    private static final class InheritableThreadLocalMap extends InheritableThreadLocal<Map<Object, Object>> {
        @Override
        protected Map<Object, Object> initialValue() {
            return Maps.newHashMap();
        }

        @SuppressWarnings("unchecked")
        @Override
        protected Map<Object, Object> childValue(Map<Object, Object> parentValue) {
            if (parentValue instanceof HashMap) {
                return (Map<Object, Object>) ((HashMap<Object, Object>) parentValue).clone();
            } else {
                return null;
            }
        }
    }

}
